#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Towards and atmosphere more favourable to firestorm development in Europe """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd

# Dew Point Temperature calculation
def Tdcalc(H,p,T):
	ep = 0.622  #g/g
	e0 = 611.3  #Pa
	b = 17.2694 
	T1 = 273.15
	T2 = 35.86  
	Td = (b*T1-T2*xr.ufuncs.log(H*p/(ep*e0)))/(b-xr.ufuncs.log(H*p/(ep*e0)))-273.15
	Td = xr.where(T<Td,T,Td)
	return Td

# K-index calculation
def Kicalc(T850,T500,Td850,T700,Td700):
	Ki = (T850-T500)+Td850-(T700-Td700)
	return Ki	

GCM=sys.argv[1]        # Enter GCM
RCM=sys.argv[2]        # Enter RCM
rcp=sys.argv[3]        # Enter RCP scenario
year=int(sys.argv[4])  # Enter year


path_hus850='./EURO-CORDEX/'+RCM+'/'+GCM+'/hus850/'  # 850 hPa specific humidity in g/g
path_hus700='./EURO-CORDEX/'+RCM+'/'+GCM+'/hus700/'  # 700 hPa specific humidity in g/g
path_ta850='./EURO-CORDEX/'+RCM+'/'+GCM+'/ta850/'    # 500 hPa temperature in Kelvin
path_ta700='./EURO-CORDEX/'+RCM+'/'+GCM+'/ta700/'    # 700 hPa temperature in Kelvin
path_ta500='./EURO-CORDEX/'+RCM+'/'+GCM+'/ta500/'    # 850 hPa temperature in Kelvin
path_lm='./EURO-CORDEX/'+RCM+'/'+GCM+'/lm/'          # Landmask

path_out_Ki='./K-index/'+RCM+'/'+GCM+'/'

lm = xr.open_dataset(path_landmask+'landmask.nc')['sftlf']

if (GCM=='EC-EARTH'):
	GCM2='ICHEC-EC-EARTH'
	dia2='31'
elif (GCM=='HadGEM2-ES'):
	GCM2='MOHC-HadGEM2-ES'
	dia2='30'
elif (GCM=='MPI-ESM-LR'):
	GCM2='MPI-M-MPI-ESM-LR'
	dia2='31'
elif (GCM=='NorESM1-M'):
	GCM2='NCC-NorESM1-M'
	dia2='31'

if ano<2006:
	experiment='historical'
else:
	experiment=rcp

## Landmask
lm = xr.open_dataset(path_landmask+'landmask.nc')['sftlf']
## 850 hPa specific humidity
hus850 = xr.open_dataset(path_hus850+'hus850_EUR-11_'+GCM2+'_'+experiment+'_day_'+str(ano)+'.nc',decode_times=False)['hus850']
## 700 hPa specific humidity
hus700 = xr.open_dataset(path_hus700+'hus700_EUR-11_'+GCM2+'_'+experiment+'_day_'+str(ano)+'.nc',decode_times=False)['hus700']
## 850 hPa air temperature
ta850 = xr.open_dataset(path_ta850+'ta850_EUR-11_'+GCM2+'_'+experiment+'_day_'+str(ano)+'.nc',decode_times=False)['ta850']-273.15
## 700 hPa air temperature
ta700 = xr.open_dataset(path_ta700+'ta700_EUR-11_'+GCM2+'_'+experiment+'_day_'+str(ano)+'.nc',decode_times=False)['ta700']-273.15
## 500 hPa air temperature
ta500 = xr.open_dataset(path_ta500+'ta500_EUR-11_'+GCM2+'_'+experiment+'_day_'+str(ano)+'.nc',decode_times=False)['ta500']-273.15

# Creating time vector
if (GCM=='EC-EARTH'):
	dr = pd.date_range(start=str(ano)+'-01-01', end=str(ano)+'-12-31', freq='1D')
	dates = dr
elif (GCM=='HadGEM2-ES'):
	dr = pd.date_range(start=str(ano)+'-01-01', end=str(ano)+'-12-31', freq='1D')
	dates = dr[((dr.day != 29) | (dr.month != 2)) & ((dr.day != 31) | (dr.month <4))]
elif (GCM=='MPI-ESM-LR'):
	dr = pd.date_range(start=str(ano)+'-01-01', end=str(ano)+'-12-31', freq='1D')
	dates = dr
elif (GCM=='NorESM1-M'):
	dr = pd.date_range(start=str(ano)+'-01-01', end=str(ano)+'-12-31', freq='1D')
	dates = dr[(dr.day != 29) | (dr.month != 2)]
hus850['time'] = dates
hus700['time'] = dates
ta850['time'] = dates
ta700['time'] = dates
ta500['time'] = dates

# Loop for computing daily K-index
Ki = xr.full_like(ta850,np.nan)
for t in range(ta850.time.size):
	print(dates[t])

	H850 = hus850.isel(time=t).copy()
	H700 = hus700.isel(time=t).copy()
	T850 = ta850.isel(time=t).copy()
	T700 = ta700.isel(time=t).copy()
	T500 = ta500.isel(time=t).copy()

	Td850 = Tdcalc(H850,85000.,T850)
	Td700 = Tdcalc(H700,70000.,T700)
	Ki.values[t,:] = Kicalc(T850,T500,Td850,T700,Td700)

print('Year: {} \n Loop end. Saving files...'.format(str(ano)))

Ki.name = 'Ki'
Ki.attrs['standard_name'] = 'Ki'
Ki.attrs['long_name'] = 'K-index'
Ki.attrs['units'] = 'Celsius'

# Saving files
Ki.to_netcdf(path_out_Ki+'Ki_'+experiment+'_'+str(ano)+'.nc')
